package com.stone.framework.batch.support;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.sql.DataSource;

import org.springframework.batch.core.partition.support.Partitioner;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.core.ColumnMapRowMapper;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.util.Assert;

public class JdbcTemplatePartitioner implements Partitioner, InitializingBean {

    private DataSource dataSource;

    private NamedParameterJdbcOperations jdbcTemplate;

    private String sql;

    private String keyField;

    @Override
    public Map<String, ExecutionContext> partition(int gridSize) {
        Map<String, ExecutionContext> map = new HashMap<String, ExecutionContext>(gridSize);

        List<Map<String, Object>> list = jdbcTemplate.query(this.sql, new ColumnMapRowMapper());
        if (list == null || list.isEmpty()) {
            return map;
        }

        for (int i = 0; i < list.size(); i++) {
            Map<String, Object> dto = list.get(i);
            ExecutionContext context = new ExecutionContext();
            dto.keySet().forEach(key -> context.put(key, dto.get(key)));
            map.put(dto.get(keyField).toString(), context);
        }
        return map;
    }

    @Override
    public void afterPropertiesSet() throws IllegalArgumentException {
        Assert.notNull(this.dataSource, "dataSource must not be null");
        this.jdbcTemplate = new NamedParameterJdbcTemplate(this.dataSource);
    }

    public void setDataSource(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    public void setSql(String sql) {
        this.sql = sql;
    }

    public void setKeyField(String keyField) {
        this.keyField = keyField;
    }

}
